#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：刘骏轩2025讯飞 
@File ：LineFind.py
@Author ：ljx
@Date ：2025/4/25 上午 9:30
'''

import time
import cv2
import numpy as np
from VideoOperator import videoOperator


def IQR(data):
    '''
    基于四分位距（IQR）去除离群值
    :param data: 一组数据
    :return:
    '''
    Q1 = np.percentile(data, 25)  # 第一四分位数
    Q3 = np.percentile(data, 75)  # 第三四分位数
    IQR = Q3 - Q1

    # 定义离群值范围（通常1.5×IQR）
    lower_bound = Q1 - 1.5 * IQR
    upper_bound = Q3 + 1.5 * IQR

    filtered_data = data[(data >= lower_bound) & (data <= upper_bound)]
    new_mean = np.mean(filtered_data)
    return new_mean


class LineFinder:
    def __init__(self):
        # 图像参数
        # 去畸变前
        self.orign_height = 480
        self.orign_width = 640
        # 去畸变后
        self.height = 377
        self.width = 409
        # 图像预处理时的裁剪
        self.cut_height = 100
        self.cut_width = 409

        # 预处理参数
        # lab 颜色提取
        self.lab_upper = np.array([255, 180, 180])
        self.lab_lower = np.array([180, 120, 120])
        # 普通的阈值处理
        self.thresh = 200

        # 中间分界黑线
        self.black = 188
        # 判断黑线的白色像素点的个数标准
        self.judgeBlackLine = 3
        # 这个参数是表示只有左右min_pixes都为黑色的线，才会是黑线
        self.black_min_pixes = 6

        # 去畸变参数
        self.mtx = np.array([[856.33915146, 0, 322.63971813], [0, 860.39225183, 253.318471], [0, 0, 1]])
        self.dist = np.array([[-9.80598534e-01, 2.57605537e+00, -4.85808691e-03, -1.26090067e-02, -6.82896742e+00]])
        self.camera_mtx = np.array([[548.00073245, 0, 278.7087777], [0, 676.68335634, 264.78544216], [0, 0, 1]])

        # self.calibration_roi=[72,66,409,377]
        # self.calibration_x,self.calibration_y,self.calibration_w,self.calibration_h = self.calibration_roi
        self.calibration_x = 72
        self.calibration_y = 66
        # self.calibration_w = 409
        # self.calibration_h = 377
        self.map_x, self.map_y = cv2.initUndistortRectifyMap(self.mtx, self.dist, None, self.camera_mtx,
                                                             (self.orign_width, self.orign_height), cv2.CV_32FC1)

        # 滑动窗口找线时的参数
        self.window_min_pixels = 0
        self.window_num = 10
        self.window_height = self.cut_height // self.window_num

        '''
        area : 车车在巡线时所处的状态：
                1：straight：正常的直线区域
                2：end：终点附近
                3、stop：终点
                4、circle：圆弧中，只有单边线，根据外边线进行判断
                5、transition：过度区域 , 在圆弧区域中，若丢失了外边线，则进入该区域
                6、transition_roll : 旋转进入直线区域
        '''
        self.area = "straight"
        # 在环岛区域中外线的位置 left or right
        self.outLine = "right"

        self.count = 0

    def calibration(self, img):
        """
        对图像去畸变
        :param img:
        :return:
        """
        dst = cv2.remap(img, self.map_x, self.map_y, cv2.INTER_LINEAR)
        dst = dst[self.calibration_y:self.calibration_y + self.height,
              self.calibration_x:self.calibration_x + self.width]
        return dst

    def do_pre(self, cut_img: np.ndarray, way=2) -> np.ndarray:
        if way == 1:
            gray = cv2.cvtColor(cut_img, cv2.COLOR_BGR2GRAY)
            # cv2.imshow('gray', gray)

            blur = cv2.GaussianBlur(gray, (5, 5), 1)
            # cv2.imshow('blur', blur)

            # 普通的阈值分割
            _, mask = cv2.threshold(blur, self.thresh, 255, cv2.THRESH_BINARY)
            # cv2.imshow('way1mask', mask)

            # kernel = np.ones((5, 5), np.uint8)
            # 开运算
            # pre_img = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

            # cv2.imshow('pre_img', mask)
        else:
            lab_img = cv2.cvtColor(cut_img, cv2.COLOR_BGR2Lab)

            # blur = cv2.GaussianBlur(lab_img, (5, 5), 1)

            mask = cv2.inRange(lab_img, self.lab_lower, self.lab_upper)
            # cv2.imshow('mask', mask)

            # kernel = np.ones((5, 5), np.uint8)

            # 开运算
            # pre_img = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

            # cv2.imshow('pre_img', pre_img)

        # 进行腐蚀操作，去除部分噪点
        # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))  # 矩形结构 或者 kernel = np.ones((5, 5), np.uint8)
        # kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))  # 椭圆结构
        # kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (5, 5))  # 十字形结构
        # 做开运算，去除外部白点,先腐蚀后膨胀
        # #进行腐蚀操作，去除部分噪点
        # kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))  # 椭圆结构
        # pre_img = cv2.erode(pre_img, kernel)  # 腐蚀
        #
        # #膨胀恢复
        # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))  # 矩形结构
        # pre_img = cv2.dilate(pre_img, kernel)
        # pre_img = cv2.morphologyEx(pre_img, cv2.MORPH_OPEN, kernel)

        # cv2.waitKey(1)

        return mask

    def getBlackLine(self, bool_img):
        '''
        得到分界的黑线
        :param bool_img: 预处理图片的bool型数组
        :return: 返回黑列的列值
        '''
        finally_list = np.sum(bool_img, axis=0)  # 沿着纵轴求和，判断一列中有多少白色点

        if finally_list[188] <= self.judgeBlackLine:
            flag = True
            '''其左右min_pixes个像素的区域内不能出现白线区域，不然认为是赛道的断裂'''
            for i in range(1, self.black_min_pixes):
                if finally_list[188 - i] > self.judgeBlackLine or finally_list[188 + i] > self.judgeBlackLine:
                    # mid不是最长黑列
                    flag = False
                    break
            if flag:
                return 188

        # 如果最长黑列在极端两侧的话，就认为已经没有最长黑列的
        for i in range(1, 188 - 20, self.black_min_pixes):
            if (188 - i) >= 0 and finally_list[188 - i] <= self.judgeBlackLine:
                flag = True
                '''其左右8个像素的区域内不能出现白线区域，不然认为是赛道的断裂'''
                for j in range(1, 8):
                    if 188 - i - j < 0:
                        flag = False
                        break
                    if (finally_list[188 - i - j] > self.judgeBlackLine or
                            finally_list[188 - i + j] > self.judgeBlackLine):
                        flag = False
                        break
                if flag:
                    return 188 - i

            if (188 + i) < self.cut_width and finally_list[188 + i] <= self.judgeBlackLine:
                flag = True
                for j in range(1, self.black_min_pixes):
                    if 188 + i + j >= self.cut_width:
                        flag = False
                        break
                    if (finally_list[188 + i - j] > self.judgeBlackLine or
                            finally_list[188 + i + j] > self.judgeBlackLine):
                        flag = False
                        break
                if flag:
                    return 188 + i

        return None

    def get_LR_Line(self, bool_img):
        """
        get left and right line
        :param bool_img: bool 图片
        :return:

        """
        '''采用滑动窗口的方法获取左右直线'''
        '''划分窗口从下到上'''
        l_line_pos = []
        r_line_pos = []

        for x in range(self.cut_height - 1, self.window_height, -self.window_height):
            left_window = bool_img[x - self.window_height - 1:x, 0:self.black]
            right_window = bool_img[x - self.window_height - 1:x, self.black:self.width]
            '''接下来是检测线的位置'''
            # 找左窗口的线位置
            _, y_vals = left_window.nonzero()
            if len(y_vals) > self.window_min_pixels:
                # print(f"线中点的个数{len(y_vals)}")
                # l_avg_y_val = IQR(y_vals)
                l_avg_y_val = y_vals.mean()
                l_line_pos.append([x, int(l_avg_y_val)])

            _, y_vals = right_window.nonzero()
            if len(y_vals) > self.window_min_pixels:
                # print(f"线中点的个数{len(y_vals)}")
                # r_avg_y_val = IQR(y_vals)
                r_avg_y_val = y_vals.mean()
                r_line_pos.append([x, self.black + int(r_avg_y_val)])

        return l_line_pos, r_line_pos

    def least_squares_fit(self, x, y):
        """
        使用最小二乘法拟合 y = kx + b
        :param x: 自变量数组
        :param y: 因变量数组
        :return: 拟合直线的系数 k, b
        """
        if len(x) < 2 or len(y) < 2:
            # print("拟合失败")
            return None, None

        x = np.array(x)
        y = np.array(y)

        n = len(x)
        x_mean = np.mean(x)
        y_mean = np.mean(y)

        # 计算 k, b
        k = np.sum((x - x_mean) * (y - y_mean)) / np.sum((x - x_mean) ** 2)
        b = y_mean - k * x_mean

        return k, b

    def get_Mid(self, l_line, r_line):
        '''

        :param l_line: 左路坐标点
        :param r_line: 右路坐标点
        :way :采用的方法：1、根据两边的点直接求出其对应的中点坐标,然后根据点拟合出中线方程 2、根据两边直线的方程求出对应的中线方程
        :return: 中线

        '''
        mid = []

        if l_line is None and r_line is None:
            print("左线右线均丢失")
            return None

        l_index = r_index = 0

        base = 6 if self.area == "circle" else 0

        while l_index < len(l_line) and r_index < len(r_line):
            if l_line[l_index][0] == r_line[r_index][0]:
                mid.append((l_line[l_index][0], int((l_line[l_index][1] + r_line[r_index][1]) / 2)))
                l_index += 1
                r_index += 1
            elif l_line[l_index][0] > r_line[r_index][0]:
                mid.append((l_line[l_index][0], int((l_line[l_index][1] + self.cut_width) / 2) - base))
                l_index += 1
                # print("右点缺失")
            else:
                mid.append((r_line[r_index][0], int((r_line[r_index][1]) / 2) + base))
                r_index += 1
                # print("左点缺失")
        while l_index < len(l_line):
            mid.append((l_line[l_index][0], int((l_line[l_index][1] + self.cut_width) / 2) - base))
            l_index += 1
            # print("右点缺失")
        while r_index < len(r_line):
            mid.append((r_line[r_index][0], int((r_line[r_index][1]) / 2) + base))
            r_index += 1
            # print("左点缺失")

        return np.array(mid)

        # else:  # 分别拟合左右两直线，然后求出中线
        #     l_y = [pt[0] for pt in l_line]  # y 是图像行号
        #     l_x = [pt[1] for pt in l_line]  # x 是图像列号
        #
        #     r_y = [pt[0] for pt in r_line]
        #     r_x = [pt[1] for pt in r_line]
        #
        #     l_k, l_b = self.least_squares_fit(l_y, l_x)
        #     r_k, r_b = self.least_squares_fit(r_y, r_x)
        #     # print(f"l_k{l_k},l_b{l_b},r_k{r_k},r_b{r_b}")
        #     k = (l_k + r_k) / 2
        #     b = (l_b + r_b) / 2
        # print(f"中线拟合结果：x = {k:.4f}y + {b:.4f}")

    def solve(self, img, way=None):
        '''
        完整的处理代码
        :param way:
        :param img:
        :return:
        '''
        calibration = self.calibration(img)

        cut_img = calibration[self.height - self.cut_height:self.height, 0:self.width]

        pre_img = self.do_pre(cut_img)
        # cv2.imwrite('/home/ucar/ucar_ws/src/vision/scripts/test/test_'+str(self.count)+'.jpg',pre_img)

        bool_img = pre_img > 100

        cv2.imshow('pre', pre_img)
        # print("找左右线")

        # 找到分界黑列
        self.black = self.getBlackLine(bool_img)
        if self.black is None:
            print("丢失黑线")
            '''
                没有最长黑线，就是前面有一条长直线
                两种情况：end 或者 过了circle后的transition过渡区域
                        end 区域前是 straight，可以通过这个进行判断
                        transition 前是circle，circle在外线区域消失后就会进入transition
            '''
            if bool_img[self.cut_height - 1].sum() > 200:
                self.area = "stop"
                print(self.area)
                return None
            else:
                self.area = "end"
                print(self.area)
                return None
        else:

            l_line, r_line = self.get_LR_Line(bool_img)

            if self.area == "straight":
                if self.outLine == "right":
                    if len(l_line) <= 0:
                        self.area = "circle"
                else:
                    if len(r_line) <= 0:
                        self.area = "circle"
                print(f"area:{self.area},outline:{self.outLine}")

            if way == "test":
                # 画出最长黑线
                cv2.line(cut_img, (self.black, 0), (self.black, self.height), (0, 0, 0), 2)
                # 中线
                cv2.line(cut_img, (210, 0), (210, self.height), (128, 0, 128), 2)
                # 标注左车道线为红色
                if len(l_line) != 0:
                    for point in l_line:
                        y, x = point  # 注意顺序是 [y, x]
                        cv2.circle(cut_img, (x, y), 3, (0, 0, 255), -1)  # 红色点
                # else:
                #     print("没有左点")
                # 标注右车道线为蓝色
                if len(r_line) != 0:
                    for point in r_line:
                        y, x = point
                        cv2.circle(cut_img, (x, y), 3, (0, 255, 0), -1)  # 绿点
                # else:
                #     print("没有右点")

            # 不选择拟合直线，直接通过中点的error来调控
            mid = self.get_Mid(l_line, r_line)

            # 展示中线位置
            if way == "test":
                if len(mid) != 0:
                    for point in mid:
                        y, x = point
                        # print("中点：",x)
                        cv2.circle(cut_img, (x, y), 3, (0, 255, 255), -1)

            if len(mid) == 0:
                print("无法中点")
                return None

            ''' 188是cut_width/2 '''
            error = (mid[:, 1].mean() - 210) * 10 / self.width

            # print(error)
            # print(f"area:{self.area}")
            # if error > 0:
            #     print("turn right")
            # else:
            #     print("turn left")

            # if way == "test":
            #     cv2.imshow('cut', cut_img)
            #     cv2.imshow('calibration', calibration)
            #     cv2.waitKey(1)

            # cv2.imwrite('/home/ucar/ucar_ws/src/vision/scripts/test/ori_' + str(self.count) + '.jpg', cut_img)

            self.count += 1

            return error


if __name__ == '__main__':
    finder = LineFinder()
    cap = videoOperator()

    while 1:
        img = cap.Get_Image()
        if img is None:
            continue
        finder.solve(img, "test")
